from typing import Dict, List
from instruction_re.formatters.format_utils import *
from instruction_re.core.datatypes import Entity
from ..instances.EntityFormatter import EntityFormatter
from ..instances.MainTaskFormatter import MainTaskFormatter


class NERTask:
    """
    Task: extract entity and their types from the sentence
    """

    def __init__(self, template_order) -> None:
        self.template_order = template_order

    @classmethod
    def format_context(cls, data_item: Dict):
        sentence = data_item["sentence"]
        return add_sentence_prefix(sentence)

    def format_data_sample(
        self,
        data_item,
        instruction: str,
        options: List[str] = None,
    ) -> Dict:
        relations = data_item["relations"]

        context = self.format_context(data_item)

        instruction = add_instruction_prefix(instruction)
        options_joined = add_options_prefix(format_str_list_to_a_string(options))

        entities = self.extract_entities(relations)
        label = EntityFormatter.format_entities(
            entities, template_order=self.template_order
        )

        return {
            "instruction": instruction + " " + options_joined,
            "context": context + " " + Preffix.ANSWER.value,
            "labels": label,
        }

    def format_entity_for_simple_template(self, entity: Entity) -> str:
        return f"{entity.name} ({entity.type})"

    @classmethod
    def extract_entities(cls, relations):
        entities = []
        for r in relations:
            entities.append(Entity.from_json(r["head"]))
            entities.append(Entity.from_json(r["tail"]))

        # 对 entities 进行去重
        entities = list(set(entities))
        return entities
